add important linear algebra operations by marcelluethi · Pull Request #125 · dimwit-dev/dimwit

marcelluethi · 2026-07-02T15:12:21Z

This PR adds important linear algebra operations, such as qr, cholesky, svd, solve and eigh.
Furthermore, it puts the implementation of the LinearAlgebraTensorOps (such as norm, det, inv) into the LinearAlgebra module and makes them just simple wrappers.

This commit wraps important linear algebra operations, such as qr, cholesky, svd, solve and eigh.

…rAlgebra module

benikm91 · 2026-07-03T11:39:27Z

+  */
+object LinearAlgebra:
+
+  enum VectorNormType:


This is wrong. Norms should be different functions, not a parameter switching a method's behavior.

Additionally, this is wrong as Ord(1) and Ord(2) are identical to L1 and L2!

Let's improve the interface, maybe:

t.norm(p=2)
t.norm(p=Double.PositiveInfinity)

or

t.pNorm(p=2)

The choice of L1 and L2 as important special cases for Ord(1) and Ord(2) was a deliberate choice. These are the two cases that a user needs most often. Ord instead is used to have arbitrary (even fractional) norms, which are less common. t.norm(VectorNormType.L2) is more readable than t.norm(VectorNormType.Ord(2))

I don't see why having an enum to select the type of norm is wrong. It is just a design choice and as discoverable and typesafe as t.normInfinity or t.normL2` I have, however, no strong opinion on this point and could be easily convinced of another design.

benikm91 · 2026-07-03T11:40:52Z

+    case Ord(p: Double)
+    case Inf
+
+  enum MatrixNormType:


See previous comment. Same error here.

Lets think of better interface, here maybe multiple functions instead of parameter changing behavior.

see above answer

benikm91 · 2026-07-03T11:47:29Z

+    *
+    * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.qr.html#jax.numpy.linalg.qr JAX documentation]] for more details on the underlying implementation.
+    */
+  def qr[L1: Label, L2: Label, V: IsFloating](t: Tensor2[L1, L2, V], mode: QRMode = QRMode.Reduced): (q: Tensor2[L1, L2, V], r: Tensor2[L1, L2, V]) =


I don't know QR, but checking documentation these types look wrong:

Same names imply to me same shapes, but Q and R have different shapes.

Should be Tensor2[L1, K, V], Tensor2[K, L2, V]
for "reduced", where K is maybe a label we define ourselves or an axis the user must provide?

Should be Tensor2[Prime[L1], L1, V], Tensor2[L1, L2, V]
for "complete"

Different return type indicates to me that these should be two different functions.

FYI: There is also a missing "raw" option (okay by me to not support full interface, yet).

This is an excellent point. We need to introduce an additional label LBasis, that is not used but denotes the label for the newly introduced basis dimension. I will add that.

I was aware of the raw option, but did not see the application for us as the moment.

benikm91 · 2026-07-03T11:52:19Z

+    * @param axis2 The second axis along which to compute the determinant.
+    * @return A new tensor with the determinant computed, where the two specified axes are removed
+    */
+  def det[T <: Tuple: Labels, L1: Label, L2: Label, V: IsFloating](t: Tensor[T, V], axis1: Axis[L1], axis2: Axis[L2])(using


A minimal interface design would require a Tensor2 as a determinant. Even a square matrix.

The type should therefore be Tensor2[L1, L2] or maybe even Tensor2[L, Prime[L]]

Should also have integers implementation, not just Floating-point

The question here is how close we want to be to the original Jax design. For me, enforcing a Rank 2 Tensor would be fine, as long as we do it consistently. I would, however, not enforce Tensor2[L,Prime[L]] as any invertible linear transformation has a determinant. I don't think introducing a SquareTensor type is worth the added complexity at this points. But we can keep it in mind for the future.

benikm91 · 2026-07-03T11:53:46Z

+    * @param offset The offset of the diagonal from the main diagonal. Positive values indicate diagonals above the main diagonal, while negative values indicate diagonals below it.
+    * @return A new tensor with the diagonal extracted, where the two specified axes are replaced by a new 1D axis labeled L1.
+    */
+  def diagonal[T <: Tuple, L1: Label, L2: Label, V](t: Tensor[T, V], axis1: Axis[L1], axis2: Axis[L2], offset: Int = 0)(using


Maybe the same. Should be defined on Tensor2. Instead of axes parameters user must use vmap.

I guess here it does not require to be squared.

benikm91 · 2026-07-03T11:55:44Z

+    *
+    * @return a new tensor with the same shape as t, but with the last two axes replaced by their inverses.
+    */
+  def inv[T <: Tuple: Labels, V: IsFloating](t: Tensor[T, V]): Tensor[T, V] = Tensor(Jax.jnp.linalg.inv(t.jaxValue))


I guess this is just wrong as it does not require Axes and it not a Tensor2. So current implementation falls back on last two axes (JAX). Should also be based on Tensor2 type. Tensor must be squared so again L1, L2 or L, Prime[L] must be decided.

Yes, it is maybe better to let that be a Tensor2. Again, it is not as general as Jax, but we can always vmap

benikm91 · 2026-07-03T11:58:27Z

+    *
+    * @return A new 0-D tensor containing the computed norm of the input tensor.
+    */
+  def norm[T <: Tuple, V: IsFloating](t: Tensor[T, V]): Tensor0[V] =


This "norms" all the dimensions by squaring all the numbers and summing them up.

I think I would not wrap/not support this. This seems to be over permissive implementation and not a mathematical / theoretical construct.

Note a user can always do "t.pow(2).sum"

I am unsure about this. it is not as easy as t.pow(2).sum, at least not if we only have vector and matrix norms as in the current implementation. We would need nested reduce implementations to compute something like this.

Norm is such a fundamental operation, that it seems useful to have a method that works without much ceremony.

This method is, how the current norm in LinearAlgebraOps is implemented.

benikm91 · 2026-07-03T12:01:35Z

+    *
+    * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.cholesky.html#jax.numpy.linalg.cholesky JAX documentation]] for more details on the underlying implementation.
+    */
+  def cholesky[L1: Label, L2: Label, V: IsFloating](t: Tensor2[L1, L2, V], upper: Boolean = false, symmetrizeInput: Boolean = true): Tensor2[L1, L2, V] =


Is keeping the same labels logical? I would have to think about it. Implementation looks fine.

What would be interesting is if we can express NaN in the value type: Float32 | NaN, as cholesky can return NaN type, but maybe lets push that into the future.

The label seem fine to me, as the cholesky factors map into the same space.

I would also push more type safety into the future.

benikm91 · 2026-07-03T12:03:33Z

+    *
+    * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.eigh.html#jax.numpy.linalg.eigh JAX documentation]] for more details on the underlying implementation.
+    */
+  def eigh[L1: Label, L2: Label, V: IsFloating](t: Tensor2[L1, L2, V], upper: Boolean = false, symmetrizeInput: Boolean = true)


Again t must be squared. So maybe L and Prime[L]? Here maybe more important than before because of eigenvalues being a vector L1 (or L2).

I think the correct solution is to introduce two more labels that the user must provide: One for the Tensor1 of Eigenvalues and one for the Tensor2 of EigenVectors. The signature could be

def eigh[L1: Label, L2: Label, LEig: Label, LSpace: Label, V: IsFloating](t: Tensor2[L1, L2, V], eigAxis: Axis[LEig], spaceAxis: Axis[LSpace], upper: Boolean = false, symmetrizeInput: Boolean = true) : (eigenvalues: Tensor1[LEig, V], eigenvectors: Tensor2[LSpace, LEig, V])

Again, I would postpone enforcing a square matrix.

benikm91 · 2026-07-03T12:07:32Z

+    * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.svd.html#jax.numpy.linalg.svd JAX documentation]] for more details on the underlying implementation.
+    */
+  def svd[L1: Label, L2: Label, V: IsFloating](t: Tensor2[L1, L2, V], fullMatrices: Boolean = false, hermitian: Boolean = false)
+      : (U: Tensor2[L1, L2, V], S: Tensor1[L1, V], Vh: Tensor2[L1, L2, V]) =


S: Tensor1[K, V] with K being either L1 or L2 based on which is bigger.

Maybe easiest to resolve by an assert statement, requiring L1 >= L2

Types are wrong if fullMatrices is true. U would have to be L1, L1 and V would be L2, L2. Maybe drop fullMatrices support for now?

benikm91 · 2026-07-03T12:10:02Z

+    *
+    * @see [[https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.solve.html#jax.numpy.linalg.solve JAX documentation]] for more details on the underlying implementation.
+    */
+  def solve[L1: Label, L2: Label, V: IsFloating](a: Tensor2[L1, L2, V], b: Tensor1[L1, V]): Tensor1[L2, V] =


a is squared so only one label type L, makes return value not ambiguous.

Note: Only supports part of solve interface (b can also be matrix in JAX), but okay for me.

Again, i think square does not mean that they have to have the same label

benikm91 · 2026-07-03T12:11:05Z

-      * @return A new tensor with the trace computed, where the two specified axes are removed, and the remaining axes are preserved.
+      * @see [[LinearAlgebra.trace]] for details
      */
    def trace[L1: Label, L2: Label](axis1: Axis[L1], axis2: Axis[L2], offset: Int = 0)(using


Okay :) I am the sinner :) I implemented it wrong here already. But we should clean it up.

Agree, let's make everything 2D

benikm91

See comments

marcelluethi · 2026-07-04T15:21:18Z

@benikm91 I addressed the label issued in Eigh, QR and SVD, but left the other issues unresolved so far, until we have a clear consensus how to proceed.

marcelluethi · 2026-07-07T04:32:43Z

Pushed a commit, which simplifies the operations by defining them only on rank 2 Tensors.
It does, however, not change anything with the norm. Still looking for a good solution there.

add important linear algebra operations

2790a19

This commit wraps important linear algebra operations, such as qr, cholesky, svd, solve and eigh.

marcelluethi requested a review from benikm91 July 2, 2026 15:12

implement LinearAlgebraOps extension methods in terms using the Linea…

c8d018e

…rAlgebra module

marcelluethi force-pushed the add-linalg-operations branch from 9f82248 to c8d018e Compare July 2, 2026 15:26

benikm91 reviewed Jul 3, 2026

View reviewed changes

benikm91 requested changes Jul 3, 2026

View reviewed changes

introduce additional axes for qr, eigh and svd

770534e

simplify linear algebra operation to work only on rank 2 tensors

bad0956

marcelluethi added 2 commits July 7, 2026 06:48

add an separate axis label LSing for the singular axis in svd

98e143b

add tests for diagonal methods

4b463e0

Uh oh!

Conversation

marcelluethi commented Jul 2, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

marcelluethi Jul 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

benikm91 Jul 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

benikm91 left a comment

Choose a reason for hiding this comment

Uh oh!

marcelluethi commented Jul 4, 2026

Uh oh!

marcelluethi commented Jul 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

marcelluethi Jul 4, 2026 •

edited

Loading

benikm91 Jul 3, 2026 •

edited

Loading